from fastapi import FastAPI, HTTPException, Request, Header
from pydantic import BaseModel, Field, validator, conint, confloat
from services.auth import verify_api_key
from services.qdrant import QdrantService
from services.dify_adapter import to_dify_format
import os
import json
import logging
import re
from typing import Optional

app = FastAPI()

# 初始化 Qdrant 客户端
qdrant = QdrantService(
    host=os.getenv("QDRANT_HOST", "192.168.109.128"),
    port=int(os.getenv("QDRANT_PORT", 6333)),
    api_key=os.getenv("QDRANT_API_KEY")
)


@app.post("/retrieval")
async def retrieval(
    request: Request,  # 必须保留此参数
    authorization: str = Header(None)
):
    # 1. 认证校验
    if not verify_api_key(authorization):
        raise HTTPException(403, {"error_code": 1001, "error_msg": "Invalid API Key"})
    
    # 2. 解析请求体（兼容中文符号和数字格式）
    try:
        raw_body = await request.body()
        body_str = raw_body.decode('utf-8')
        
        # 修复常见JSON语法错误：中文逗号、数字引号
        body_str = body_str.replace('，', ',')  # 中文逗号 → 英文逗号
        body_str = re.sub(r'"(\d+)"', r'\1', body_str)  # 移除数字引号，如 "5" → 5
        
        payload = json.loads(body_str)
        print(f"接收到请求体: {payload}")  # 调试输出
    except json.JSONDecodeError:
        logging.error(f"无效JSON: {body_str}")
        raise HTTPException(400, detail={
            "error_code": 1002, 
            "error_msg": "JSON格式错误：请检查逗号/引号用法"
        })
    
    # 3. 提取参数并设置默认值
    knowledge_id = payload.get("knowledge_id", "")
    query = payload.get("query", "")
    top_k = payload.get("retrieval_setting", {}).get("top_k", 5)
    score_threshold = payload.get("retrieval_setting", {}).get("score_threshold", 0.7)

    # 打印knowledge_id和query、top_k、score_threshold
    print(f"Knowledge ID: {knowledge_id}, Query: {query}, Top K: {top_k}, Score Threshold: {score_threshold}")

    if(top_k < 40):
        # logging.warning(f"top_k值过小: {top_k}，重置为40")
        print(f"top_k值过小: {top_k}，重置为40")
        top_k = 40
    if(score_threshold == 0.0  ):
        # logging.warning(f"score_threshold值为0，重置为0.01")
        print(f"score_threshold值为0，重置为0.01")
        score_threshold = 0.01
    
    if score_threshold > 0.1:
        # logging.warning(f"score_threshold值过小: {score_threshold}，重置为0.1")
        print(f"score_threshold值过大: {score_threshold}，重置为0.01")
        score_threshold = 0.01

    # 4. 参数验证
    if not knowledge_id or not query:
        raise HTTPException(400, detail={
            "error_code": 1003,
            "error_msg": "必填参数缺失：knowledge_id 和 query 不能为空"
        })

    # 5. 执行向量检索（含异常处理）
    try:
        results = await qdrant.search(
            collection_name=knowledge_id,
            query_text=query,
            limit=top_k,
            score_threshold=score_threshold
        )
        return {"records": to_dify_format(results)}
    except Exception as e:
        # 打印异常e的错误堆栈信息
        logging.exception("检索过程中发生错误", e)
        # 记录错误信息
        logging.error(f"检索失败: {str(e)}")
        raise HTTPException(500, detail={
            "error_code": 1004,
            "error_msg": f"Qdrant服务错误：{str(e)}"
        })

@app.get("/health")
def health_check():
    try:
        qdrant_status = qdrant.ping()
    except Exception as e:
        logging.error(f"Qdrant连接检查失败: {str(e)}")
        qdrant_status = False
    return {"status": "alive", "qdrant_connected": qdrant_status}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=9000)